#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/11/13 23:06
# @Author  : huidong.bai
# @File    : ECNRSession.py
# @Software: PyCharm
# @Mail    : MasterBai2018@outlook.com

import os
from ctypes import *
from conftest import logging
from src.utils.common import is_available_file
import re
import pdb


class PNRApiCode:
    PNR_API_CODE_OK = 0
    PNR_API_CODE_INVALID_SAMPLE_RATE = -1
    PNR_API_CODE_INVALID_QBIT = -2
    PNR_API_CODE_INVALID_PARAM = -3
    PNR_API_CODE_INVALID_WORKMODE = -4
    PNR_API_CODE_UNSUPPORT_OPERATION = -5
    PNR_API_CODE_NO_MEMORY = -6
    PNR_API_CODE_INVALID_ENGINE = -7


class PNRWorkMode:
    mode_VR = 0
    mode_BT = 1
    mode_CARPLAY = 2
    mode_WECHAT = 3
    mode_CARLINK = 4
    mode_HICAR = 5
    mode_CARLINK_VR = 6
    mode_HICAR_VR = 7
    mode_WECHAT_VR = 8
    mode_COMMON = 9
    mode_DOWNLINK = 10


class PNRConfigST(Structure):
    _fields_ = [("type", c_char * 32),
                ("mic_num", c_int),
                ("mic_distance", c_int),
                ("channel_num", c_int)]


class AIBSPNRDataST(Structure):
    _fields_ = [("input", POINTER(c_char)),
                ("output", POINTER(c_char))]


class ECNRSession:
    m_engine = None
    m_library = None
    m_config = None
    m_synthMethod = None
    m_outputChannelNum = None

    def __init__(self, lib_path: str):
        try:
            self.m_library = cdll.LoadLibrary(lib_path)
        except Exception as e:
            logging.error(f"Load pnr library error: {e}")
            pass
    
    def output_channel(self, number):
        binary_representation = bin(number)[2:]
        count_of_ones = binary_representation.count('1')
        return count_of_ones

    def create_pnr_engine(self, device_type: str, mic_num: int, channel_num: int, mic_distance: int):
        if not device_type or mic_num <= 0 or channel_num < mic_num or mic_distance <= 0:
            logging.error(f"Create pnr engine error, device_type:{device_type}, mic_num:{mic_num}, channel_num:{channel_num}, mic_distance:{mic_distance}")
            return -1
        self.m_config = PNRConfigST()
        self.m_config.type = device_type.encode()
        self.m_config.mic_num = mic_num
        self.m_config.channel_num = channel_num
        self.m_config.mic_distance = mic_distance
        self.m_engine = self.m_library.create_pnr_engine(pointer(self.m_config))

        if self.m_engine is None:
            logging.error(f'Create pnr engine failed, the engine is:{self.m_engine}')
            return -2
        return 0

    def create_linein_pnr_engine(self, device_type: str, mic_num: int, channel_num: int, mic_distance: int):
        if not device_type or mic_num <= 0 or channel_num < mic_num or mic_distance <= 0:
            logging.error(f"Create linein pnr engine error, device_type:{device_type}, mic_num:{mic_num}, channel_num:{channel_num}, mic_distance:{mic_distance}")
            return -1
        self.m_config = PNRConfigST()
        self.m_config.type = device_type.encode()
        self.m_config.mic_num = mic_num
        self.m_config.channel_num = channel_num
        self.m_config.mic_distance = mic_distance
        self.m_engine = self.m_library.create_linein_pnr_engine(pointer(self.m_config))

        if self.m_engine is None:
            logging.error(f'Create linein pnr engine failed, the engine is:{self.m_engine}')
            return -2
        return 0

    def start_pnr_engine(self, input_sample_rate, output_sample_rate, audio_format, output_channel_mask, synth_method):
        self.m_synthMethod = synth_method
        self.m_outputChannelNum = self.output_channel(output_channel_mask)
        self.m_library.start_pnr_engine.argtypes = [c_void_p, c_int, c_int, c_int, c_int, c_int]
        self.m_library.start_pnr_engine.restype = c_int
        ret = self.m_library.start_pnr_engine(self.m_engine, input_sample_rate, output_sample_rate, audio_format, output_channel_mask, synth_method)
        if ret != 0:
            logging.error(f'Start pnr engine failed, ret:{ret}')
            return ret
        return 0

    def start_linein_pnr_engine(self, input_sample_rate, output_sample_rate, audio_format, output_channel_mask, synth_method):
        self.m_synthMethod = synth_method
        self.m_outputChannelNum = self.output_channel(output_channel_mask)
        self.m_library.start_pnr_engine.argtypes = [c_void_p, c_int, c_int, c_int, c_int, c_int]
        self.m_library.start_pnr_engine.restype = c_int
        ret = self.m_library.start_pnr_engine(self.m_engine, input_sample_rate, output_sample_rate, audio_format, output_channel_mask, synth_method)
        if ret != 0:
            logging.error(f'Start linein pnr engine failed, ret:{ret}')
            return ret
        return 0
    
    def process_linein_data(self, audio_path: str, sample_rate: int, output_path: str):
        self.m_library.process_pnr_data.argtypes = [c_void_p, POINTER(AIBSPNRDataST), c_int, c_long]
        self.m_library.process_pnr_data.restype = c_int
        if not audio_path or not sample_rate:
            logging.error(f'Bad audio or rate, audio_path:{audio_path}, sample_rate:{sample_rate}')
            return -1

        t_frame = self.get_pnr_frame_size()
        channel_frame_length = sample_rate * 2 // 1000 * t_frame
        output_length = channel_frame_length * self.m_outputChannelNum
        frame_size = self.m_config.channel_num * channel_frame_length

        audio_length = os.path.getsize(audio_path)
        # 循环读取，并送音频
        with open(audio_path, 'rb') as fp:
            # 获取文件头信息
            header = fp.read(100)
            if header[:4] == b"RIFF":
                # 如果是wav文件，则获取头信息中的音频长度
                pos = re.search(b'data', header).end()
                audio_length = int.from_bytes(header[pos:pos + 4], byteorder='little')
                # 将文件指针移动到音频数据的起始位置
                fp.seek(pos + 4, 0)
            else:
                fp.seek(0, 0)

            timestamp = 0
            output_file = open(output_path, "wb")
            # 实际pcm音频流数据
            remaining_bytes = audio_length
            while remaining_bytes > 0:
                if remaining_bytes >= frame_size:
                    buffer = fp.read(frame_size)
                else:
                    buffer = fp.read(remaining_bytes)

                nr_data = AIBSPNRDataST()
                nr_data.input = cast(create_string_buffer(buffer), POINTER(c_char))
                nr_data.output = cast(create_string_buffer(b'\x00' * output_length), POINTER(c_char))

                timestamp += 10
                ret = self.m_library.process_linein_data(self.m_engine, pointer(nr_data), t_frame, timestamp)

                if ret == 0:
                    output_file.write(nr_data.output[:output_length])
                else:
                    output_file.close()
                    logging.error("process linein nr data failed or the current task don't need the rest of data.")
                    print("process linein nr data failed or the current task don't need the rest of data.")
                    return -1

                remaining_bytes -= len(buffer)
                # 如果已经读取了指定长度，则退出循环
                if remaining_bytes <= 0:
                    break
            output_file.close()
        return 0

    def process_pnr_data(self, audio_path: str, sample_rate: int, output_path: str):
        self.m_library.process_pnr_data.argtypes = [c_void_p, POINTER(AIBSPNRDataST), c_int, c_long]
        self.m_library.process_pnr_data.restype = c_int
        if not audio_path or not sample_rate:
            logging.error(f'Bad audio or rate, audio_path:{audio_path}, sample_rate:{sample_rate}')
            return -1

        t_frame = self.get_pnr_frame_size()
        channel_frame_length = sample_rate * 2 // 1000 * t_frame
        output_length = channel_frame_length * self.m_outputChannelNum
        print("m_synthMethod:", self.m_synthMethod)
        print("m_outputChannelNum:", self.m_outputChannelNum)
        print("回调长度: ", output_length)
        frame_size = self.m_config.channel_num * channel_frame_length

        audio_length = os.path.getsize(audio_path)
        # 循环读取，并送音频
        with open(audio_path, 'rb') as fp:
            # 获取文件头信息
            header = fp.read(100)
            if header[:4] == b"RIFF":
                # 如果是wav文件，则获取头信息中的音频长度
                pos = re.search(b'data', header).end()
                audio_length = int.from_bytes(header[pos:pos + 4], byteorder='little')
                # 将文件指针移动到音频数据的起始位置
                fp.seek(pos + 4, 0)
            else:
                fp.seek(0, 0)

            timestamp = 0
            output_file = open(output_path, "wb")
            # 实际pcm音频流数据
            remaining_bytes = audio_length
            while remaining_bytes > 0:
                if remaining_bytes >= frame_size:
                    buffer = fp.read(frame_size)
                else:
                    buffer = fp.read(remaining_bytes)

                nr_data = AIBSPNRDataST()
                nr_data.input = cast(create_string_buffer(buffer), POINTER(c_char))
                nr_data.output = cast(create_string_buffer(b'\x00' * output_length), POINTER(c_char))

                timestamp += 10
                ret = self.m_library.process_pnr_data(self.m_engine, pointer(nr_data), t_frame, timestamp)

                if ret == 0:
                    offset = 0
                    channel_list = []
                    for i in range(self.m_outputChannelNum):
                        channel_list.append(nr_data.output[offset:offset+channel_frame_length])
                        offset += channel_frame_length
                    
                    offset_bytes = 0
                    for i in range(160):
                        for data in channel_list:
                            output_file.write(data[offset_bytes:offset_bytes+2])
                        offset_bytes += 2

                else:
                    output_file.close()
                    logging.error("process nr data failed or the current task don't need the rest of data.")
                    print("process nr data failed or the current task don't need the rest of data.")
                    return -1

                remaining_bytes -= len(buffer)
                # 如果已经读取了指定长度，则退出循环
                if remaining_bytes <= 0:
                    break
            output_file.close()
        return 0

    def free_linein_pnr_engine(self):
        self.m_library.free_linein_pnr_engine.argtypes = [c_void_p]
        self.m_library.free_linein_pnr_engine.restype = c_int
        return self.m_library.free_linein_pnr_engine(self.m_engine)

    def free_pnr_engine(self):
        self.m_library.free_pnr_engine.argtypes = [c_void_p]
        self.m_library.free_pnr_engine.restype = c_int
        return self.m_library.free_pnr_engine(self.m_engine)

    def set_pnr_work_mode(self, mode, sample_rate):
        self.m_library.set_pnr_work_mode.argtypes = [c_void_p, c_int, c_int]
        self.m_library.set_pnr_work_mode.restype = c_int
        return self.m_library.set_pnr_work_mode(self.m_engine, c_int(mode), c_int(sample_rate))

    def set_linein_pnr_work_mode(self, mode, sample_rate, spectrum):
        self.m_library.set_linein_pnr_work_mode.argtypes = [c_void_p, c_int, c_int]
        self.m_library.set_linein_pnr_work_mode.restype = c_int
        return self.m_library.set_linein_pnr_work_mode(self.m_engine, c_int(mode), c_int(sample_rate), c_int(spectrum))

    def set_pnr_mic_mute_option(self, option):
        self.m_library.set_pnr_mic_mute_option.argtypes = [c_void_p, c_bool]
        self.m_library.set_pnr_mic_mute_option.restype = c_int
        return self.m_library.set_pnr_mic_mute_option(self.m_engine, option)

    def set_pnr_enable_option(self, option):
        self.m_library.set_pnr_enable_option.argtypes = [c_void_p, c_bool]
        self.m_library.set_pnr_enable_option.restype = c_int
        return self.m_library.set_pnr_enable_option(self.m_engine, option)

    def set_pnr_audio_quality(self, quality):
        self.m_library.set_pnr_audio_quality.argtypes = [c_void_p, c_ushort]
        self.m_library.set_pnr_audio_quality.restype = c_int
        return self.m_library.set_pnr_audio_quality(self.m_engine, quality)

    def get_pnr_version(self):
        self.m_library.get_pnr_version.argtypes = [c_void_p]
        self.m_library.get_pnr_version.restype = c_char_p
        return self.m_library.get_pnr_version(self.m_engine)

    def get_pnr_HFT_param(self):
        self.m_library.get_pnr_HFT_param.argtypes = [c_void_p]
        self.m_library.get_pnr_HFT_param.restype = c_char_p
        return self.m_library.get_pnr_HFT_param(self.m_engine)

    def get_pnr_MVR_param(self):
        self.m_library.get_pnr_MVR_param.argtypes = [c_void_p]
        self.m_library.get_pnr_MVR_param.restype = c_char_p
        return self.m_library.get_pnr_MVR_param(self.m_engine)

    def get_pnr_Gen_param(self):
        self.m_library.get_pnr_Gen_param.argtypes = [c_void_p]
        self.m_library.get_pnr_Gen_param.restype = c_char_p
        return self.m_library.get_pnr_Gen_param(self.m_engine)

    def get_pnr_frame_size(self):
        self.m_library.get_pnr_frame_size.argtypes = [c_void_p]
        self.m_library.get_pnr_frame_size.restype = c_int
        return self.m_library.get_pnr_frame_size(self.m_engine)

    def get_linein_pnr_frame_size(self):
        self.m_library.get_linein_pnr_frame_size.argtypes = [c_void_p]
        self.m_library.get_linein_pnr_frame_size.restype = c_int
        return self.m_library.get_linein_pnr_frame_size(self.m_engine)


if __name__ == '__main__':
    pnr_engine = ECNRSession(lib_path="/data1/MasterBai/mango/lib/libNR_dynamic.so")
    pnr_engine.create_pnr_engine(device_type="lexus_2S", mic_num=2, channel_num=9, mic_distance=646)

    # 开始音频处理
    ret = pnr_engine.start_pnr_engine(input_sample_rate=16000, output_sample_rate=16000, audio_format=16, output_channel_mask=1, synth_method=1);
    if ret != 0:
        print("start pnr engine failed, code:%d", ret)
        pnr_engine.free_pnr_engine();

    pnr_engine.process_pnr_data(audio_path="/data1/MasterBai/mango/audio/pnr_2_7_9.wav", sample_rate=16000, output_path="/data1/MasterBai/mango/output2222222.pcm")
